import argparse
import subprocess
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

def run_editor():
    """运行障碍物编辑器"""
    print("Starting obstacle editor...")
    try:
        subprocess.run([sys.executable, "editor/obstacle_editor.py"], check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error running editor: {e}")

def run_training(agent_type="sb3"):
    """运行训练脚本"""
    print(f"Starting training with {agent_type} agent...")
    try:
        subprocess.run([sys.executable, "train.py", "--agent", agent_type], check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error during training: {e}")

def run_game(mode="human", model_path=None, use_sb3=False):
    """运行游戏"""
    print(f"Starting game in {mode} mode...")
    try:
        cmd = [sys.executable, "play_game.py", "--mode", mode]
        if model_path:
            cmd.extend(["--model", model_path])
        if use_sb3:
            cmd.append("--sb3")
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error running game: {e}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="Ball Battle Game Runner")
    subparsers = parser.add_subparsers(dest="command", help="Available commands")
    
    # 编辑器子命令
    editor_parser = subparsers.add_parser("editor", help="Run obstacle editor")
    
    # 训练子命令
    train_parser = subparsers.add_parser("train", help="Run training")
    train_parser.add_argument("--agent", choices=["custom", "sb3"], default="sb3",
                              help="Choose agent type (default: sb3)")
    
    # 游戏子命令
    game_parser = subparsers.add_parser("play", help="Play the game")
    game_parser.add_argument("--mode", choices=["human", "ai"], default="human",
                             help="Choose game mode (default: human)")
    game_parser.add_argument("--model", type=str, help="Path to trained model")
    game_parser.add_argument("--sb3", action="store_true", help="Use Stable-Baselines3 model")
    
    # 安装依赖子命令
    install_parser = subparsers.add_parser("install", help="Install dependencies")
    
    args = parser.parse_args()
    
    if args.command == "editor":
        run_editor()
    elif args.command == "train":
        run_training(args.agent)
    elif args.command == "play":
        run_game(args.mode, args.model, args.sb3)
    elif args.command == "install":
        print("Installing dependencies...")
        try:
            subprocess.run([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], check=True)
            print("Dependencies installed successfully!")
        except subprocess.CalledProcessError as e:
            print(f"Error installing dependencies: {e}")
    else:
        # 默认显示帮助信息
        parser.print_help()

if __name__ == "__main__":
    main()